import pandas as pd   
import numpy as np   
   
def Log(sr):
    #自然对数函数
    return np.log(sr)

def Rank(sr):
    #列-升序排序并转化成百分比
    return sr.expanding(min_periods=2).rank(pct=True)

def Delta(sr,period):
    #period日差分
    return sr.diff(period)

def Delay(sr,period):
    #period阶滞后项
    return sr.shift(period)

def Corr(x,y,window):
    #window日滚动相关系数
    t = x.rolling(window, min_periods = 1).corr(y)
    t.replace([np.inf, -np.inf], np.nan, inplace = True)
    t = t.fillna(0)
    return t
    
def Cov(x,y,window):
    #window日滚动协方差
    return x.rolling(window, min_periods = 1).cov(y)

def Sum(sr,window):
    #window日滚动求和
    return sr.rolling(window, min_periods = 1).sum()

def Prod(sr,window):
    #window日滚动求乘积
    #return df.rolling(window, min_periods = 1)).apply(lambda x: np.prod(x))
    sr1 = sr.copy()
    for i in range(window-1):
        sr1 = sr1*sr.shift(i+1)
    return sr1    

def Mean(sr,window):
    #window日滚动求均值
    return sr.rolling(window, min_periods = 1).mean()

def Std(sr,window):
    #window日滚动求标准差
    return sr.rolling(window, min_periods = 1).std()

def Tsrank(sr, window):
    #window日序列末尾值的顺位
    return sr.rolling(window, min_periods=1).rank(pct=True)

               
def Tsmax(sr, window):
    #window日滚动求最大值    
    return sr.rolling(window, min_periods = 1).max()

def Tsmin(sr, window):
    #window日滚动求最小值    
    return sr.rolling(window, min_periods = 1).min()

def Sign(sr):
    #符号函数
    return np.sign(sr)

def Max(sr,n):
    #比较取大函数
    return np.maximum(n, sr)

def Max2(sr1,sr2):
    #比较取小函数
    return np.maximum(sr1, sr2)

def Min(sr,n):
    #比较取小函数
    return np.minimum(n, sr)

def Min2(sr1,sr2):
    #比较取小函数
    return np.minimum(sr1, sr2)


def Sma(sr,n,m):
    #sma均值
    #df1 = df.ewm(alpha=m/n).mean()
    return sr.ewm(alpha=m/n, adjust=False).mean()

def Abs(sr):
    #求绝对值
    return sr.abs()

def Sequence(n):
    #生成 1~n 的等差序列
    #return np.arange(1,n+1)
    return pd.Series(np.arange(1,n+1).tolist())


def rolling_regression_b1(y1):
    #注意这个0对应的是'HS300_return'列系数
    x1 = np.arange(1,(len(y1)+1))
    #注意这个1对应的是'指数y_return'列系数
    beta1=(y1*(x1-x1.mean())).sum()/((x1-x1.mean())*(x1-x1.mean())).sum() #从线性回归公式的角度计算beta1
    return beta1

def Regbeta(sr,B,window):
    df = sr
    t = df.rolling(window,  min_periods=3).apply(rolling_regression_b1, raw=True, engine="numba")
    t.replace([np.inf, -np.inf], np.nan, inplace = True)
    t = t.fillna(0)
    return t

def rolling_regressionAB_b1(x):
    # print(x.shape)
    x1=x[:,0] #注意这个0对应的是'HS300_return'列系数
    y1=x[:,1] #注意这个1对应的是'指数y_return'列系数
    beta1=(y1*(x1-x1.mean())).sum()/(((x1-x1.mean())*(x1-x1.mean())).sum()+0.000000001) #从线性回归公式的角度计算beta1
    return beta1

def RegbetaAB(x, y, n):
    # try:
        df = pd.concat([x ,y], axis=1)
        t = df.rolling(n, method="table", min_periods=3).apply(rolling_regressionAB_b1, raw=True, engine="numba")
        t.replace([np.inf, -np.inf], np.nan, inplace=True)
        t = t.fillna(0)
        return t.iloc[:,0]

def gen_dy(x):
    weights = np.arange(1,len(x)+1,1)
    y = weights / weights.sum()
    return (x*y).sum()

def Decaylinear(sr, window):  #将dataframe运算转成np数组运算
    weights = np.arange(1,window+1,1)
    sr1 = sr.rolling(window, min_periods=1).apply(gen_dy, raw=True, engine="numba")
    return sr1

def Lowday(sr,window):
    #计算sr前window期时间序列中最小值距离当前时点的间隔
    sr1 = sr.rolling(window, min_periods=1, closed='left').apply(np.argmin,raw=True, engine="numba") + 1
    return window - sr1


def Highday(sr,window):
    #计算sr前window期时间序列中最大值距离当前时点的间隔
    sr1 = sr.rolling(window, min_periods=1, closed='left').apply(np.argmax,raw=True, engine="numba") + 1
    return window - sr1


def gen_y(x):
    weights = 0.9*np.arange(len(x)-1,0-1,-1)
    return (x*weights).sum()


def Wma(sr,window):   
    sr1 = sr.rolling(window, min_periods=1).apply(gen_y , raw=True, engine="numba")
    return sr1

def Count(part,window):
    #计算前n期满足条件condition的样本个数,此时输入的part为0、1变量
    return part.rolling(window, min_periods=1, closed='left').sum()


def Sumif(part,window):
    #对前n项条件求和，part为条件筛选后的数据
    return part.rolling(window, min_periods=1, closed='left').sum()
